#include "SenGenerator.h"
#include "Param.h"
#include "Utils.h"
#include <vector>
#include "Default.h"
using namespace std;

void SenGenerator::genDeprel(CGRNN* net, Matrix* rep, int k, vector< pair<real,int> >& ret) {
	Param* params = net->params;
	Matrix* score = new Matrix(params->bdr->length);
	Matrix::gemv(1, CblasNoTrans, params->Wdr, rep, 0, score);
	score->addi(params->bdr);
	real* prob = new real[score->length];
	Utils::safelyComputeSoftmax(prob, score->data, score->length);
	Utils::kmax(prob, score->length, k, ret);
	delete score; delete[] prob;
}

void SenGenerator::genPos(CGRNN* net, Matrix* rep, int deprel, int k, vector< pair<real,int> > &ret) {
	Param* params = net->params;
	Matrix* score = Matrix::zeros(params->bpos->length);
	Matrix::gemv(1, CblasNoTrans, params->Wpos, rep, 0, score);
	Matrix* temp = params->Wdrpos->getColumn(deprel);
	score->addi(temp)->addi(params->bpos);
	real* prob = new real[score->length];
	Utils::safelyComputeSoftmax(prob, score->data, score->length);
	Utils::kmax(prob, score->length, k, ret);
	delete score; delete temp; delete[] prob;
}

void SenGenerator::genWord(CGRNN* net, Matrix* rep, int pos, int deprel, int k, vector< pair<real,int> > &ret) {
	Param* params = net->params;
	Dictionary* vocaDic = net->vocaDic;

	Matrix* score = Matrix::zeros(vocaDic->size());
	Matrix* temp = Matrix::zeros(params->Wword.size());

	for (int i = 0; i < temp->length; i++) {
		temp->data[i] = Matrix::dot(rep, params->Wword[i]);
	}

	for (int word = 0; word < vocaDic->size(); word++) {
		int* path = vocaDic->path[word];
		real* code = vocaDic->code[word];
		int codeLen = vocaDic->codeLength[word];
		for (int i = 0; i < codeLen; i++) {
			real x = temp->data[path[i]] + params->Wdrword->get(path[i], deprel)
							+ params->Wposword->get(path[i], pos) + params->bword->data[path[i]];
			x *= code[i];
			score->data[word] += log(1./(1+exp(-x)));
		}
	}

	Utils::kmax(score->data, score->length, k, ret);
	delete score; delete temp;
}

real SenGenerator::lengthPenalty(int length, int expectedLength) {
	return -abs(length - expectedLength)*2;
	//return Utils::logPoissonProb(length, expectedLength);
}

#define GEN_LEFT 0
#define GEN_RIGHT 1
#define GEN_DONE 2

void SenGenerator::genLeftDep(CGRNN* net, Deptree* tree, real score, STATE *state, AGENDA& ret, int expectedLength) {
	int headnode = state->headnode;
	int lastnode = state->lastnode;
	int genstate = state->genstate;

	Param* params = net->params;
	int dim = net->dim;

	if (genstate == GEN_LEFT) {
		Matrix* priorRep = state->priorRep->dup();
		Matrix* inputHead_Left = state->inputHead_Left;
		Matrix* inputHead_Right = state->inputHead_Right;
		Matrix* inputHead_Left_p = state->inputHead_Left_p;
		Matrix* inputHead_Right_p = state->inputHead_Right_p;

		Matrix* pGRep;

		Matrix* input = new Matrix(dim);
		Matrix::gemv(1, CblasNoTrans, params->Wgs_p[0], priorRep, 0, input);
		input->addi(inputHead_Left_p);
		pGRep = net->func(input); delete input;

		// generate deprel
		vector<pair<real,int>> drCands;
		genDeprel(net, pGRep, kDeprels, drCands);

		for (pair<real,int> drcand : drCands) {
			int deprel = drcand.second;

			if (deprel == params->bdr->length-1) { // EOC
				Deptree* newtree = tree->clone(true);
				Matrix* newpriorRep = priorRep->dup();

				Matrix* newinputHead_Left = inputHead_Left->dup();
				Matrix* newinputHead_Right = new Matrix(net->dim);
				Matrix::gemv(1, CblasNoTrans, params->Wgh[1], newtree->headRep[headnode], 0, newinputHead_Right);
				Matrix::gemv(1, CblasNoTrans, params->Wgp[1], newtree->gRep[headnode], 1, newinputHead_Right);
				newinputHead_Right->addi(params->bg[1]);

				Matrix* newinputHead_Left_p = inputHead_Left_p->dup();
				Matrix* newinputHead_Right_p = new Matrix(net->dim);
				Matrix::gemv(1, CblasNoTrans, params->Wgh_p[1], newtree->headRep[headnode], 0, newinputHead_Right_p);
				Matrix::gemv(1, CblasNoTrans, params->Wgp_p[1], newtree->gRep[headnode], 1, newinputHead_Right_p);
				newinputHead_Right_p->addi(params->bg_p[1]);

				vector<Matrix*> newhiProds;
				for (Matrix* m : state->hiProds)
					newhiProds.push_back(m->dup());

				real newscore = score + log(drcand.first);
				ret[newscore] = pair<Deptree*,STATE*>(newtree,
						new STATE(headnode, lastnode, GEN_RIGHT, newinputHead_Left, newinputHead_Right,
								newinputHead_Left_p, newinputHead_Right_p, newpriorRep, newhiProds));
			}
			// generate pos / word / cap
			else if (lastnode < tree->nNodes -1){
				vector<pair<real,int>> posCands;
				genPos(net, pGRep, deprel, kPos, posCands);
				
				for (pair<real,int> poscan : posCands) {
					int pos = poscan.second;

					vector<pair<real,int>> wordCands;
					genWord(net, pGRep, pos, deprel, kWords, wordCands);

					for (pair<real,int> wrdcan : wordCands) {
						int word = wrdcan.second;
						int newlastnode = lastnode + 1;

						Deptree* newtree = tree->clone(true);
						newtree->deprel[newlastnode] = deprel;
						newtree->pos[newlastnode] = pos;
						newtree->word[newlastnode] = word;
						newtree->dir[newlastnode] = 0;
						newtree->parent[newlastnode] = headnode;

						Matrix* temp = new Matrix(net->dim);
						Matrix* Li = params->L[word];
						Matrix::gemv(1, CblasNoTrans, params->Ww, Li, 0, temp);
						temp->addi(params->Lpos[pos])->addi(params->bw);
						newtree->headRep[newlastnode] = net->func(temp);
						delete temp;

						newtree->children[headnode][0][newtree->nChildren[headnode][0]] = newlastnode;
						newtree->nChildren[headnode][0]++;

						Matrix* input = new Matrix(net->dim);
						Matrix::gemv(1, CblasNoTrans, params->Wg_p[0][deprel], newtree->headRep[newlastnode], 0, input);
						Matrix::gemv(1, CblasNoTrans, params->Wtr[0], priorRep, 1, input);
						input->addi(params->btr[0]);
						Matrix* newpriorRep = net->func(input); delete input;


						Matrix* newinputHead_Left = inputHead_Left->dup();
						Matrix* newinputHead_Right = inputHead_Right->dup();

						Matrix* newinputHead_Left_p = inputHead_Left_p->dup();
						Matrix* newinputHead_Right_p = inputHead_Right_p->dup();

						vector<Matrix*> newhiProds;
						for (Matrix* m : state->hiProds)
							newhiProds.push_back(m->dup());
						Matrix* prod = new Matrix(dim);
						Matrix::gemv(1, CblasNoTrans, params->Wg[0][deprel], newtree->headRep[newlastnode], 0, prod);
						newhiProds.push_back(prod);

						real newscore = score - lengthPenalty(lastnode+1,expectedLength)
							+ log(drcand.first) + log(poscan.first) + wrdcan.first
							+ lengthPenalty(newlastnode+1,expectedLength);

						ret[newscore] = pair<Deptree*,STATE*>(newtree,
								new STATE(headnode, newlastnode, GEN_LEFT,
										newinputHead_Left, newinputHead_Right, newinputHead_Left_p, newinputHead_Right_p,
										newpriorRep, newhiProds));
					}
				}
			}
		}

		delete priorRep;
		delete pGRep;
	}
}

void SenGenerator::genRightDep(CGRNN* net, Deptree* tree, real score, STATE *state, AGENDA& ret, int expectedLength) {
	int headnode = state->headnode;
	int lastnode = state->lastnode;
	int genstate = state->genstate;

	Param* params = net->params;

	if (genstate == GEN_RIGHT) {
		Matrix* priorRep = state->priorRep->dup();
		Matrix* inputHead_Left = state->inputHead_Left;
		Matrix* inputHead_Right = state->inputHead_Right;
		Matrix* inputHead_Left_p = state->inputHead_Left_p;
		Matrix* inputHead_Right_p = state->inputHead_Right_p;

		Matrix* pGRep;

		Matrix* input = new Matrix(dim);
		Matrix::gemv(1, CblasNoTrans, params->Wgs_p[1], priorRep, 0, input);
		input->addi(inputHead_Right_p);
		pGRep = net->func(input); delete input;

		// generate deprel
		vector<pair<real,int>> drCands;
		genDeprel(net, pGRep, kDeprels, drCands);

		for (pair<real,int> drcand : drCands) {
			int deprel = drcand.second;

			if (deprel == params->bdr->length-1) { // EOC
				Deptree* newtree = tree->clone(true);
				real newscore = score + log(drcand.first);
				
				// compute gRep for children
				Matrix* in = Matrix::zeros(net->dim);
				for (int j = newtree->nChildren[headnode][0] + newtree->nChildren[headnode][1] - 1; j >= 0; j--) {
					int i = lastnode - j;
					int parent = headnode;
					int dir = newtree->dir[i];
					Matrix* inputParent = dir == 0 ? inputHead_Left : inputHead_Right;

					int nSisters = newtree->nChildren[parent][0] + newtree->nChildren[parent][1] - 1;
					if (nSisters == 0) {
						Matrix* temp = inputParent->add(params->anonRep[dir]);
						newtree->gRep[i] = net->func(temp); delete temp;
					}
					else {
						in->fill(0);
						for (int sis = 0; sis < (int)state->hiProds.size(); sis++) {
							if (sis + lastnode - newtree->nChildren[parent][0] + newtree->nChildren[parent][1] + 1 != i)
								in->addi(state->hiProds[sis]);
						}
						in->divi(CGRNN_NORM ? nSisters : 1)->addi(inputParent);
						newtree->gRep[i] = net->func(in);
					}
				}
				delete in;

				// turn to new headnode
				int newheadnode = headnode + 1;
				if (newheadnode <= newtree->nNodes-1 && newtree->headRep[newheadnode] != NULL) {
					Matrix* newpriorRep = params->anonRep[0]->dup();

					Matrix* newinputHead_Left = Matrix::zeros(net->dim);
					Matrix::gemv(1, CblasNoTrans, params->Wgh[0], newtree->headRep[newheadnode], 0, newinputHead_Left);
					Matrix::gemv(1, CblasNoTrans, params->Wgp[0], newtree->gRep[newheadnode], 1, newinputHead_Left);
					newinputHead_Left->addi(params->bg[0]);
					Matrix* newinputHead_Right = Matrix::zeros(net->dim);

					Matrix* newinputHead_Left_p = Matrix::zeros(net->dim);
					Matrix::gemv(1, CblasNoTrans, params->Wgh_p[0], newtree->headRep[newheadnode], 0, newinputHead_Left_p);
					Matrix::gemv(1, CblasNoTrans, params->Wgp_p[0], newtree->gRep[newheadnode], 1, newinputHead_Left_p);
					newinputHead_Left_p->addi(params->bg_p[0]);
					Matrix* newinputHead_Right_p = Matrix::zeros(net->dim);

					vector<Matrix*> newhiProds;
					ret[newscore] = pair<Deptree*,STATE*>(newtree,
							new STATE(newheadnode, lastnode, GEN_LEFT, newinputHead_Left, newinputHead_Right,
									newinputHead_Left_p, newinputHead_Right_p,
									newpriorRep, newhiProds));
				}
				else {
					vector<Matrix*> newhiProds;
					ret[newscore] = pair<Deptree*,STATE*>(newtree, new STATE(newheadnode, lastnode, GEN_DONE,
							NULL, NULL, NULL, NULL, NULL, newhiProds));
				}
			}
			// generate pos / word / cap
			else if (lastnode < tree->nNodes -1) {
				vector<pair<real,int>> posCands;
				genPos(net, pGRep, deprel, kPos, posCands);
				
				for (pair<real,int> poscan : posCands) {
					int pos = poscan.second;

					vector<pair<real,int>> wordCands;
					genWord(net, pGRep, pos, deprel, kWords, wordCands);

					for (pair<real,int> wrdcan : wordCands) {
						int word = wrdcan.second;
						
						int newlastnode = lastnode + 1;

						Deptree* newtree = tree->clone(true);
						newtree->deprel[newlastnode] = deprel;
						newtree->pos[newlastnode] = pos;
						newtree->word[newlastnode] = word;
						newtree->dir[newlastnode] = 1;
						newtree->parent[newlastnode] = headnode;

						Matrix* temp = new Matrix(net->dim);
						Matrix* Li = params->L[word];
						Matrix::gemv(1, CblasNoTrans, params->Ww, Li, 0, temp);
						temp->addi(params->Lpos[pos])->addi(params->bw);
						newtree->headRep[newlastnode] = net->func(temp);

						newtree->children[headnode][1][newtree->nChildren[headnode][1]] = newlastnode;
						newtree->nChildren[headnode][1]++;
						delete temp;

						Matrix* input = new Matrix(dim);
						Matrix::gemv(1, CblasNoTrans, params->Wg_p[1][deprel], newtree->headRep[newlastnode], 0, input);
						Matrix::gemv(1, CblasNoTrans, params->Wtr[1], priorRep, 1, input);
						input->addi(params->btr[1]);
						Matrix* newpriorRep = net->func(input); delete input;

						Matrix* newinputHead_Left = inputHead_Left->dup();
						Matrix* newinputHead_Right = inputHead_Right->dup();

						Matrix* newinputHead_Left_p = inputHead_Left_p->dup();
						Matrix* newinputHead_Right_p = inputHead_Right_p->dup();

						vector<Matrix*> newhiProds;
						for (Matrix* m : state->hiProds)
							newhiProds.push_back(m->dup());
						Matrix* prod = new Matrix(dim);
						Matrix::gemv(1, CblasNoTrans, params->Wg[1][deprel], newtree->headRep[newlastnode], 0, prod);
						newhiProds.push_back(prod);

						real newscore = score - lengthPenalty(lastnode+1,expectedLength)
							+ log(drcand.first) + log(poscan.first) + wrdcan.first
							+ lengthPenalty(newlastnode+1,expectedLength);

						ret[newscore] = pair<Deptree*,STATE*>(newtree,
								new STATE(headnode, newlastnode, GEN_RIGHT,
										newinputHead_Left, newinputHead_Right,
										newinputHead_Left_p, newinputHead_Right_p,
										newpriorRep, newhiProds));
					}
				}
			}
		}

		delete priorRep;
		delete pGRep;
	}
}

bool SenGenerator::genDep(CGRNN* net, Deptree* tree, real score, STATE *state, AGENDA& ret, int expectedLength) {
	int genstate = state->genstate;
	switch (genstate) {
	case GEN_LEFT: 
		genLeftDep(net, tree, score, state, ret, expectedLength);
		return true;
		break;
	case GEN_RIGHT:
		genRightDep(net, tree, score, state, ret, expectedLength);
		return true;
		break;
	case GEN_DONE:
		{
			vector<Matrix*> newhiProds;
			ret[score] = pair<Deptree*,STATE*>(tree->clone(),
					new STATE(state->headnode, state->lastnode, GEN_DONE, NULL, NULL, NULL, NULL, NULL, newhiProds));
		}
		return false;
		break;
	default:
		Utils::error("sthing wrong");
	}
	return false;
}

void SenGenerator::generate(CGRNN* net, Matrix* rootRep, int expectedLength, vector<pair<real,Deptree*>> &result) {
	Deptree *tree = new Deptree(maxGenLen);
	Param* params = net->params;

	tree->headRep = new Matrix*[tree->nNodes]();
	tree->gRep = new Matrix*[tree->nNodes](); // full gen. rep.
	tree->headRep[0] = Matrix::zeros(dim);
	tree->gRep[0] = rootRep->dup();

	real score = lengthPenalty(0, expectedLength);
	int headnode = 0;
	int lastnode = 0;
	vector<Matrix*> hiProds;
	Matrix* priorRep = params->anonRep[0]->dup();
	
	Matrix* inputHead_Left = new Matrix(net->dim);
	Matrix::gemv(1, CblasNoTrans, params->Wgh[0], tree->headRep[0], 0, inputHead_Left);
	Matrix::gemv(1, CblasNoTrans, params->Wgp[0], tree->gRep[0], 1, inputHead_Left);
	inputHead_Left->addi(params->bg[0]);

	Matrix* inputHead_Right = new Matrix(net->dim);
	Matrix::gemv(1, CblasNoTrans, params->Wgh[1], tree->headRep[0], 0, inputHead_Right);
	Matrix::gemv(1, CblasNoTrans, params->Wgp[1], tree->gRep[0], 1, inputHead_Right);
	inputHead_Right->addi(params->bg[1]);

	Matrix* inputHead_Left_p = new Matrix(net->dim);
	Matrix::gemv(1, CblasNoTrans, params->Wgh_p[0], tree->headRep[0], 0, inputHead_Left_p);
	Matrix::gemv(1, CblasNoTrans, params->Wgp_p[0], tree->gRep[0], 1, inputHead_Left_p);
	inputHead_Left_p->addi(params->bg_p[0]);

	Matrix* inputHead_Right_p = new Matrix(net->dim);
	Matrix::gemv(1, CblasNoTrans, params->Wgh_p[1], tree->headRep[0], 0, inputHead_Right_p);
	Matrix::gemv(1, CblasNoTrans, params->Wgp_p[1], tree->gRep[0], 1, inputHead_Right_p);
	inputHead_Right_p->addi(params->bg_p[1]);

	AGENDA agenda;
	agenda[score] = pair<Deptree*,STATE*>(tree,
			new STATE(headnode, lastnode, GEN_LEFT, inputHead_Left, inputHead_Right,
					inputHead_Left_p, inputHead_Right_p,
					priorRep, hiProds));

	bool cont = true;

	while (cont) {
		AGENDA newagenda;
		cont = false;

		//cout << agenda.size() << endl;

		for (AGENDA::iterator it = agenda.begin(); it != agenda.end(); ++it) {
			Deptree* t = it->second.first;
			STATE* tstate = it->second.second;
			real tscore = it->first;
			cont = genDep(net, t, tscore, tstate, newagenda, expectedLength) || cont;

			Deptree::free(t);
			delete tstate;
		}

		while (newagenda.size() > (unsigned int)beamsize) {
			map<real,pair<Deptree*,STATE*>>::iterator endit = newagenda.end();
			--endit;
			Deptree::free(endit->second.first);
			delete endit->second.second;
			newagenda.erase(endit);
		}

		agenda.clear();
		agenda = newagenda;
	}

	result.clear();
	for (AGENDA::iterator it = agenda.begin(); it != agenda.end(); ++it) {
		result.push_back(pair<real,Deptree*>(it->first,it->second.first));
	};
}
